{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3df823b-3b2b-4e72-9d0b-2d12f9f6297a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.train import ModelTrainer\n",
    "from sagemaker.train.configs import InputData\n",
    "from sagemaker.core.processing import (\n",
    "    ScriptProcessor,\n",
    ")\n",
    "from sagemaker.core.shapes import (\n",
    "    ProcessingInput,\n",
    "    ProcessingS3Input,\n",
    "    ProcessingOutput,\n",
    "    ProcessingS3Output\n",
    ")\n",
    "from sagemaker.serve.model_builder import ModelBuilder\n",
    "\n",
    "from sagemaker.core.workflow.parameters import (\n",
    "    ParameterInteger,\n",
    "    ParameterString,\n",
    ")\n",
    "from sagemaker.mlops.workflow.pipeline import Pipeline\n",
    "from sagemaker.mlops.workflow.steps import ProcessingStep, TrainingStep, CacheConfig\n",
    "from sagemaker.mlops.workflow.model_step import ModelStep\n",
    "from sagemaker.core.workflow.pipeline_context import PipelineSession\n",
    "from sagemaker.core.helper.session_helper import Session, get_execution_role\n",
    "from sagemaker.core import image_uris\n",
    "from sagemaker.train.configs import Compute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87016224-8d4d-4368-822c-d5b0c1febe10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the SageMaker Session\n",
    "\n",
    "sagemaker_session = Session()\n",
    "pipeline_session = PipelineSession()\n",
    "sm_client = sagemaker_session.sagemaker_client\n",
    "region = sagemaker_session.boto_region_name\n",
    "prefix = \"pipeline-v3\"\n",
    "\n",
    "account_id = sagemaker_session.account_id()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0773b04-9d54-4a24-bc10-35fb2e7dcd5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define variables and parameters needed for the Pipeline steps\n",
    "\n",
    "role = get_execution_role()\n",
    "default_bucket = sagemaker_session.default_bucket()\n",
    "base_job_prefix = \"v3-pipeline-example\"\n",
    "s3_prefix = \"v3-test-pipeline\"\n",
    "default_bucket_prefix = sagemaker_session.default_bucket_prefix\n",
    "\n",
    "# If a default bucket prefix is specified, append it to the s3 path\n",
    "if default_bucket_prefix:\n",
    "    s3_prefix = f\"{default_bucket_prefix}/{s3_prefix}\"\n",
    "    base_job_prefix = f\"{default_bucket_prefix}/{base_job_prefix}\"\n",
    "\n",
    "processing_instance_count = ParameterInteger(name=\"ProcessingInstanceCount\", default_value=1)\n",
    "training_instance_type = ParameterString(name=\"TrainingInstanceType\", default_value=\"ml.m5.xlarge\")\n",
    "model_approval_status = ParameterString(\n",
    "    name=\"ModelApprovalStatus\", default_value=\"PendingManualApproval\"\n",
    ")\n",
    "input_data = ParameterString(\n",
    "    name=\"InputDataUrl\",\n",
    "    default_value=f\"s3://sagemaker-sample-files/datasets/tabular/uci_abalone/abalone.csv\",\n",
    ")\n",
    "model_approval_status = ParameterString(\n",
    "    name=\"ModelApprovalStatus\", default_value=\"PendingManualApproval\"\n",
    ")\n",
    "\n",
    "# Cache Pipeline steps to reduce execution time on subsequent executions\n",
    "cache_config = CacheConfig(enable_caching=True, expire_after=\"30d\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3e62722-38d5-4a17-bb64-b17134b6a21a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bd85bb3-bc51-405c-8680-9dd916cc50bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile code/preprocess.py\n",
    "\n",
    "\"\"\"Feature engineers the abalone dataset.\"\"\"\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "import pathlib\n",
    "import requests\n",
    "import tempfile\n",
    "\n",
    "import boto3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.preprocessing import StandardScaler, OneHotEncoder\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler())\n",
    "\n",
    "\n",
    "# Since we get a headerless CSV file we specify the column names here.\n",
    "feature_columns_names = [\n",
    "    \"sex\",\n",
    "    \"length\",\n",
    "    \"diameter\",\n",
    "    \"height\",\n",
    "    \"whole_weight\",\n",
    "    \"shucked_weight\",\n",
    "    \"viscera_weight\",\n",
    "    \"shell_weight\",\n",
    "]\n",
    "label_column = \"rings\"\n",
    "\n",
    "feature_columns_dtype = {\n",
    "    \"sex\": str,\n",
    "    \"length\": np.float64,\n",
    "    \"diameter\": np.float64,\n",
    "    \"height\": np.float64,\n",
    "    \"whole_weight\": np.float64,\n",
    "    \"shucked_weight\": np.float64,\n",
    "    \"viscera_weight\": np.float64,\n",
    "    \"shell_weight\": np.float64,\n",
    "}\n",
    "label_column_dtype = {\"rings\": np.float64}\n",
    "\n",
    "\n",
    "def merge_two_dicts(x, y):\n",
    "    \"\"\"Merges two dicts, returning a new copy.\"\"\"\n",
    "    z = x.copy()\n",
    "    z.update(y)\n",
    "    return z\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    logger.debug(\"Starting preprocessing.\")\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--input-data\", type=str, required=True)\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    base_dir = \"/opt/ml/processing\"\n",
    "    pathlib.Path(f\"{base_dir}/data\").mkdir(parents=True, exist_ok=True)\n",
    "    input_data = args.input_data\n",
    "    bucket = input_data.split(\"/\")[2]\n",
    "    key = \"/\".join(input_data.split(\"/\")[3:])\n",
    "\n",
    "    logger.info(\"Downloading data from bucket: %s, key: %s\", bucket, key)\n",
    "    fn = f\"{base_dir}/data/abalone-dataset.csv\"\n",
    "    s3 = boto3.resource(\"s3\")\n",
    "    s3.Bucket(bucket).download_file(key, fn)\n",
    "\n",
    "    logger.debug(\"Reading downloaded data.\")\n",
    "    df = pd.read_csv(\n",
    "        fn,\n",
    "        header=None,\n",
    "        names=feature_columns_names + [label_column],\n",
    "        dtype=merge_two_dicts(feature_columns_dtype, label_column_dtype),\n",
    "    )\n",
    "    os.unlink(fn)\n",
    "\n",
    "    logger.debug(\"Defining transformers.\")\n",
    "    numeric_features = list(feature_columns_names)\n",
    "    numeric_features.remove(\"sex\")\n",
    "    numeric_transformer = Pipeline(\n",
    "        steps=[\n",
    "            (\"imputer\", SimpleImputer(strategy=\"median\")),\n",
    "            (\"scaler\", StandardScaler()),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    categorical_features = [\"sex\"]\n",
    "    categorical_transformer = Pipeline(\n",
    "        steps=[\n",
    "            (\"imputer\", SimpleImputer(strategy=\"constant\", fill_value=\"missing\")),\n",
    "            (\"onehot\", OneHotEncoder(handle_unknown=\"ignore\")),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    preprocess = ColumnTransformer(\n",
    "        transformers=[\n",
    "            (\"num\", numeric_transformer, numeric_features),\n",
    "            (\"cat\", categorical_transformer, categorical_features),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    logger.info(\"Applying transforms.\")\n",
    "    y = df.pop(\"rings\")\n",
    "    X_pre = preprocess.fit_transform(df)\n",
    "    y_pre = y.to_numpy().reshape(len(y), 1)\n",
    "\n",
    "    X = np.concatenate((y_pre, X_pre), axis=1)\n",
    "\n",
    "    logger.info(\"Splitting %d rows of data into train, validation, test datasets.\", len(X))\n",
    "    np.random.shuffle(X)\n",
    "    train, validation, test = np.split(X, [int(0.7 * len(X)), int(0.85 * len(X))])\n",
    "\n",
    "    logger.info(\"Writing out datasets to %s.\", base_dir)\n",
    "    pd.DataFrame(train).to_csv(f\"{base_dir}/train/train.csv\", header=False, index=False)\n",
    "    pd.DataFrame(validation).to_csv(\n",
    "        f\"{base_dir}/validation/validation.csv\", header=False, index=False\n",
    "    )\n",
    "    pd.DataFrame(test).to_csv(f\"{base_dir}/test/test.csv\", header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60381d4-7e7b-4016-875e-433aba80a7e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using modern ScriptProcessor class instead of SKLearnProcessor\n",
    "sklearn_pipeline_processor = ScriptProcessor(\n",
    "    image_uri=image_uris.retrieve(\n",
    "        framework=\"sklearn\",\n",
    "        region=region,\n",
    "        version=\"1.2-1\",\n",
    "        py_version=\"py3\",\n",
    "        instance_type=\"ml.m5.xlarge\",\n",
    "    ),\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    "    instance_count=processing_instance_count,\n",
    "    base_job_name=f\"{base_job_prefix}-sklearn-preprocess-job\",\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "pipe_processor_args = sklearn_pipeline_processor.run(\n",
    "    inputs=[\n",
    "        ProcessingInput(\n",
    "            input_name=\"input-1\",\n",
    "            s3_input=ProcessingS3Input(\n",
    "                s3_uri=input_data,\n",
    "                local_path=\"/opt/ml/processing/input\",\n",
    "                s3_data_type=\"S3Prefix\",\n",
    "                s3_input_mode=\"File\",\n",
    "                s3_data_distribution_type=\"ShardedByS3Key\",\n",
    "            )\n",
    "        )\n",
    "    ],\n",
    "    outputs=[\n",
    "        ProcessingOutput(\n",
    "            output_name=\"train\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=f\"s3://{sagemaker_session.default_bucket()}/{prefix}/train\",\n",
    "                local_path=\"/opt/ml/processing/train\",\n",
    "                s3_upload_mode=\"EndOfJob\"\n",
    "            )\n",
    "        ),\n",
    "        ProcessingOutput(\n",
    "            output_name=\"validation\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=f\"s3://{sagemaker_session.default_bucket()}/{prefix}/validation\",\n",
    "                local_path=\"/opt/ml/processing/validation\",\n",
    "                s3_upload_mode=\"EndOfJob\"\n",
    "            )\n",
    "        ),\n",
    "        ProcessingOutput(\n",
    "            output_name=\"test\",\n",
    "            s3_output=ProcessingS3Output(\n",
    "                s3_uri=f\"s3://{sagemaker_session.default_bucket()}/{prefix}/test\",\n",
    "                local_path=\"/opt/ml/processing/test\",\n",
    "                s3_upload_mode=\"EndOfJob\"\n",
    "            )\n",
    "        ),\n",
    "    ],\n",
    "    code=\"code/preprocess.py\",\n",
    "    arguments=[\"--input-data\", input_data],\n",
    ")\n",
    "\n",
    "step_process = ProcessingStep(\n",
    "    name=\"PreprocessAbaloneData\",\n",
    "    step_args=pipe_processor_args,\n",
    "    cache_config=cache_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9aff89f6",
   "metadata": {},
   "source": [
    "### Training Step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892bd7e6-dfbf-4434-b39e-bedac8fb209c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the output path for the model artifacts from the training job\n",
    "model_path = f\"s3://{default_bucket}/{base_job_prefix}/V3PipelineTrain\"\n",
    "\n",
    "image_uri = image_uris.retrieve(\n",
    "    framework=\"xgboost\",\n",
    "    region=region,\n",
    "    version=\"1.0-1\",\n",
    "    py_version=\"py3\",\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    ")\n",
    "\n",
    "# Using modern ModelTrainer class instead of Estimator\n",
    "model_trainer = ModelTrainer(\n",
    "    training_image=image_uri,\n",
    "    compute=Compute(\n",
    "        instance_type=\"ml.m5.xlarge\",\n",
    "        instance_count=1,\n",
    "    ),\n",
    "    base_job_name=f\"{base_job_prefix}-xgboost-train\",\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    "    hyperparameters={\n",
    "        \"objective\": \"reg:linear\",\n",
    "        \"num_round\": 50,\n",
    "        \"max_depth\": 5,\n",
    "        \"eta\": 0.2,\n",
    "        \"gamma\": 4,\n",
    "        \"min_child_weight\": 6,\n",
    "        \"subsample\": 0.7,\n",
    "        \"silent\": 0,\n",
    "    },\n",
    "    input_data_config=[\n",
    "        InputData(\n",
    "            channel_name=\"train\",\n",
    "            data_source=step_process.properties.ProcessingOutputConfig.Outputs[\"train\"].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\"\n",
    "        ),\n",
    "        InputData(\n",
    "            channel_name=\"validation\",\n",
    "            data_source=step_process.properties.ProcessingOutputConfig.Outputs[\"validation\"].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\"\n",
    "        ),\n",
    "    ],\n",
    ")\n",
    "\n",
    "train_args = model_trainer.train()\n",
    "\n",
    "step_train = TrainingStep(\n",
    "    name=\"TrainModelTestStep\",\n",
    "    step_args=train_args,\n",
    "    cache_config=cache_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab972b7",
   "metadata": {},
   "source": [
    "### Create model step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba880c73",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Model using ModelBuilder\n",
    "model_builder = ModelBuilder(\n",
    "    s3_model_data_url=step_train.properties.ModelArtifacts.S3ModelArtifacts,\n",
    "    image_uri=image_uri,\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role_arn=role,\n",
    ")\n",
    "\n",
    "step_create_model = ModelStep(\n",
    "    name=\"CreateModel\",\n",
    "    step_args=model_builder.build()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82d6d9ea",
   "metadata": {},
   "source": [
    "### Register Step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e112822",
   "metadata": {},
   "outputs": [],
   "source": [
    "step_register_model = ModelStep(\n",
    "    name=\"RegisterModel\",\n",
    "    step_args=model_builder.register(\n",
    "        model_package_group_name=\"my-model-package-group\",\n",
    "        content_types=[\"application/json\"],\n",
    "        response_types=[\"application/json\"],\n",
    "        inference_instances=[\"ml.m5.xlarge\"],\n",
    "        approval_status=\"Approved\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "764f32a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = Pipeline(\n",
    "    name=\"pipeline-v3\",\n",
    "    parameters=[\n",
    "        processing_instance_count,\n",
    "        input_data,\n",
    "        model_approval_status,\n",
    "    ],\n",
    "    steps=[step_process, step_train, step_create_model, step_register_model],\n",
    "    sagemaker_session=pipeline_session,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396c1411",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.definition()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8434a02",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline.upsert(role_arn=role)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d2d61c",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution = pipeline.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c46a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution.wait()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
